import numpy as np
from requirement import get_requirement

np.random.seed(1)


class Env():
    def __init__(self, code_service_dict, service_code_dict, qos_needed, service_qos_dict):
        self.code_service_dict = code_service_dict # code_id: [service_id]
        self.index_code_dict = dict(zip(range(len(self.code_service_dict)), sorted(self.code_service_dict.keys()))) # index: code_id
        self.code_index_dict = dict(zip(sorted(self.code_service_dict.keys()), range(len(self.code_service_dict)))) # code_id: index
        self.service_code_dict = service_code_dict # service_id: [code_id]
        self.qos_needed = qos_needed # qos约束
        self.service_qos_dict = service_qos_dict # service_id: [qos]
        self.end = 2 ** len(self.code_service_dict) - 1

   
    def get_actions(self):
        state_action_space = {}
        for i in range(2**len(self.code_service_dict)-1):
            for j in range(len(self.code_service_dict)):
                if i & 2**j == 0:
                    break
            state_action_space[i] = self.code_service_dict[self.index_code_dict[j]]
        state_action_space[2**len(self.code_service_dict)-1] = []
        return state_action_space
    
    def get_action_code_qos(self):
        action_code_qos = {}
        for service in self.service_qos_dict:
            codes = self.service_code_dict[service]
            code_state = sum(map(lambda x: 2**self.code_index_dict[x], codes))
            action_code_qos[service] = [code_state] + list(self.service_qos_dict[service])
        return action_code_qos

    def mask_to_state(self, mask):
        state = int(mask[0], 2)
        return (state, ) + mask[1:]

    def state_to_mask(self, state):
        num = len(self.code_service_dict)
        mask = '{1:0>{0}}'.format(num, bin(state[0])[2:])
        return (mask, ) + state[1:]

    def reset(self):
        return (0, 1, 1) # 已选编号，qos情况

    def step(self, state, action, weights):
        codes = self.service_code_dict[action]
        indexes = list(map(lambda x: self.code_index_dict[x], codes))
        qos = self.service_qos_dict[action] # action的 [价格，时间，评分，地点]
        code_state, qos_state = state[0], state[1:]
        
        next_code_state = code_state | sum(map(lambda x: 2**x, indexes))
        next_qos_state = [qos_state[0] - qos[0], min(qos_state[1], 1-qos[1])]
        next_state = tuple([next_code_state] + next_qos_state)
        
        for item in next_qos_state:
            if item < 0:
                return next_state, -1, True
        
        if next_state[0] == self.end:
            reward = 0
            for i in range(len(weights)):
                reward += weights[i] * next_qos_state[i]
            return next_state, reward, True
        else:
            return next_state, 0, False
        


if __name__ == '__main__':
    code_service_dict, service_code_dict, qos_needed, service_qos_dict = get_requirement()
    env = Env(code_service_dict, service_code_dict, qos_needed, service_qos_dict)